package com.cl.spark.nodenew;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.cl.spark.base.BaseSparkNode;
import com.cl.spark.dto.SparkParam;
import com.cl.spark.dto.SparkResult;
import com.cl.spark.enums.SparkNodeEnum;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
public class UnionNode extends BaseSparkNode {

    @Override
    public SparkResult process(SparkParam sparkParam) {
        List<JSONArray> dataList = sparkParam.getDataList();
        Dataset<Row> allDataset = sparkUtil.jsonArrayToDataset(dataList.get(0));
        for (int i = 1; i < dataList.size(); i++) {
            Dataset<Row> rowDataset = sparkUtil.jsonArrayToDataset(dataList.get(i));
            allDataset = allDataset.unionAll(rowDataset);
        }
        JSONArray res = new JSONArray();
        allDataset.collectAsList().forEach(row -> {
            JSONObject item = JSONObject.parseObject(row.json());
            res.add(item);
        });
        return SparkResult.success(res);
    }

    @Override
    public SparkNodeEnum getType() {
        return SparkNodeEnum.UNION;
    }
}
